# -*- encoding: utf-8 -*-
"""
========================================
@Time   ：2021/9/7 11:01
@Auther ：shutao
@File   ：utils.py
@IDE    ：PyCharm
@Github ：https://github.com/NameLacker
@Gitee  ：https://gitee.com/nameLacker
========================================
"""

import matplotlib.pyplot as plt
import numpy as np


def circle_data_point_generator(Ntrain, Ntest, boundary_gap, seed_data):
    """
    圆形决策边界两分类数据集生成器
    Args:
        Ntrain: 训练集大小
        Ntest: 测试集大小
        boundary_gap: 取值于 (0, 0.5), 两类别之间的差距
        seed_data: 随机种子

    Returns:
        'Ntrain': 训练集
        'Ntest': 测试集
    """
    train_x, train_y = [], []
    num_samples, seed_para = 0, 0
    while num_samples < Ntrain + Ntest:
        np.random.seed((seed_data + 10) * 1000 + seed_para + num_samples)
        data_point = np.random.rand(2) * 2 - 1

        # 如果数据点的模小于(0.7 - gap)，标为0
        if np.linalg.norm(data_point) < 0.7 - boundary_gap / 2:
            train_x.append(data_point)
            train_y.append(0.)
            num_samples += 1

        # 如果数据点的模大于(0.7 + gap)，标为1
        elif np.linalg.norm(data_point) > 0.7 + boundary_gap / 2:
            train_x.append(data_point)
            train_y.append(1.)
            num_samples += 1
        else:
            seed_para += 1

    train_x = np.array(train_x).astype("float64")
    train_y = np.array([train_y]).astype("float64").T

    print("训练集的维度大小 x {} 和 y {}".format(np.shape(train_x[0:Ntrain]), np.shape(train_y[0:Ntrain])))
    print("测试集的维度大小 x {} 和 y {}".format(np.shape(train_x[Ntrain:]), np.shape(train_y[Ntrain:])), "\n")

    return train_x[0:Ntrain], train_y[0:Ntrain], train_x[Ntrain:], train_y[Ntrain:]


def data_point_plot(data, label):
    """
    用以可视化生成的数据集
    Args:
        data: 形状为 [M, 2], 代表 M 2-D 数据点
        label: 取值 0 或者 1

    Returns:
        画出数据点
    """
    dim_samples, dim_useless = np.shape(data)
    plt.figure(1)
    for i in range(dim_samples):
        if label[i] == 0:
            plt.plot(data[i][0], data[i][1], color="r", marker="o")
        elif label[i] == 1:
            plt.plot(data[i][0], data[i][1], color="b", marker="o")
    plt.show()
